"""--------------------------------------------------------------------------------------------------
Filename:   viewplot.py
Author:     Chris Chester adapted from Ben Han (uha10@uclive.ac.nz)
Date:       30 August 2012
Description:
    Main plotting module for biofeedback application.
    Contains the GUI layer and plotting functions.
--------------------------------------------------------------------------------------------------"""

import matplotlib
# The recommended way to use wx with mpl is with the WXAgg backend. 
matplotlib.interactive(False)
matplotlib.use('WXAgg')

import numpy as np
import pylab
import wx, traceback
from targetfactory import TargetFactory
from decimal import Decimal
from plotpanel import PlotControl, PlotPanel
import time


#===============================================================================
# CLASS: LivePlot(wx.Frame)
#    Real-time data plotting class.
#    This class owns the data arrays containing the EMG and time values for
#    the plot points.
#===============================================================================
class LivePlot(wx.Frame):
    def __init__(self, app, parent=None, id=-1, title="Live plot frame"):
        wx.Frame.__init__(self, parent, id, title, size=(800, 600))
        self.app = app
        self.top_panel = wx.Panel(self, -1)        # top-level panel for frame
        self.plotpanel = PlotPanel(parent=self.top_panel, color=None, dpi=80, id=-1)
        self.plotctrl = PlotControl(parent=self.top_panel, id=-1, app=app)
        self.start_time = 0
        self.__do_layout()
        self.__setup_plot()
        
    def __do_layout(self):
        ''' Sets layout for Liveplot GUI '''
        sizer_frame = wx.BoxSizer(wx.VERTICAL)  # sizer for frame
        sizer_main = wx.BoxSizer(wx.VERTICAL)
        sizer_main.Add(self.plotpanel, 1, wx.EXPAND | wx.ALL, 0)
        sizer_main.Add(self.plotctrl, 0, wx.EXPAND | wx.ALL, 25)
        self.top_panel.SetSizer(sizer_main)
        sizer_frame.Add(self.top_panel, 1, wx.EXPAND, 0)
        self.SetSizer(sizer_frame)
        sizer_frame.Fit(self)
        self.ShowFullScreen(not self.IsFullScreen(), wx.FULLSCREEN_ALL)
        self.Layout()
        
    def __setup_plot(self):
        ''' Sets up the impedance plot'''
        
        # Instantiate a matplotlib axes object
        self.axes = self.plotpanel.figure.add_subplot(111)    # self.axes is an instance of an Axes object
        self.axes.set_axis_bgcolor(self.app.cfg.plotBackColour)
        self.axes.set_title('BIoimped level', size=20)
        self.axes.set_xlabel('Time (s)')
        self.axes.set_ylabel('Pressure')
        self.axes.set_autoscale_on(False)
        self.redo_yauto = True
        
        self.xaxis_max = 10

        pylab.setp(self.axes.get_xticklabels(), fontsize=13)
        pylab.setp(self.axes.get_yticklabels(), fontsize=13)
        
        self.axes.set_xbound(lower=0, upper=self.app.cfg._x_axis_max)
        self.axes.set_ybound(lower=0, upper=self.app.cfg._y_axis_max)

        
        # Initialise data arrays to be added to later
        self.app.data_lock.acquire()    # ACQUIRE DATA_LOCK <--
        self.ch1_x_data = np.array([])      # array of ch1 x values for the plot
        self.ch1_y_data = np.array([])      # array of ch1 y values for the plot
        self.ch2_x_data = np.array([])      # array of ch2 x values for the plot
        self.ch2_y_data = np.array([])      # array of ch2 y values for the plot
        
        # Used for saving all of the data points
        self.ch1_x_data_total = np.array([])     
        self.ch1_y_data_total = np.array([])      
        self.ch2_x_data_total = np.array([])      
        self.ch2_y_data_total = np.array([])      
        
        self.emgdata = []               # used to store raw EMG data
        self.plot_state = []            # used to store state information about the plot for every time point
        self.app.data_lock.release()    # RELEASE DATA_LOCK -->
        
        self.sample_period = Decimal(1) / Decimal(str(self.app.cfg.SAMPLE_RATE))
        
        # Plot the data and grab a reference to the line added
        self.emgplot = self.axes.plot(self.ch1_x_data,
                                      self.ch1_y_data,
                                      'g',
                                      self.ch2_x_data,
                                      self.ch2_y_data,
                                      'b',
                                      animated=True)
                
        self.init_plot()
        
    def __draw_impedance(self, init=False):
        ''' Draws a plot of the impedance levels '''
        self.__remove_target()


        # reset the axes
        self.axes.set_xbound(lower=0, upper=self.app.cfg._x_axis_max)
        self.axes.set_ybound(lower=0, upper=self.app.cfg._y_axis_max)
        
    def store_data(self):
        '''
        Stores the data to be used at a later time
        '''
        self.ch1_x_data_total = np.append(self.ch1_x_data_total, self.ch1_x_data) 
        self.ch1_y_data_total = np.append(self.ch1_y_data_total, self.ch1_y_data) 
        self.ch2_x_data_total = np.append(self.ch2_x_data_total, self.ch2_x_data) 
        self.ch2_y_data_total = np.append(self.ch2_y_data_total, self.ch2_y_data) 
        
    def renew_xaxis(self,emglevel1):
        ''' Refreshes the xaxis and stores the data previously plotted 
            Inputs:
            emglevel1: the emg level of channel 1(voltage from channel)
        '''
        
        s_state = self.app.sess_status
        
        # Store the previous values
        self.store_data()
        
        # Refresh the values to be plotted                    
        self.ch1_x_data = np.array([])
        self.ch1_y_data = np.array([])
        self.ch2_x_data = np.array([])
        self.ch2_y_data = np.array([])
        
        self.axes.set_xbound(lower=self.xaxis_max, upper=self.xaxis_max + self.app.cfg._x_axis_max)
        self.xaxis_max = self.xaxis_max + self.app.cfg._x_axis_max
        
        s_state.peak = (s_state.ch1_y_index, emglevel1) # New value for yaxis auto
        self.getBackground = True #New axis values means the background has to be redone
        self.redo_yauto = True        
        
    def update_data(self, emglevel1, emglevel2):
        ''' This function is called by the serial thread to append EMG values to the data array.
         Any work to be done by the serial thread from this function should 
         be added to the app.job_queue object rather than performed directly.
         Inputs:
         emglevel1: the emg level of channel 1(voltage from channel)
         emglevel2: the emg level of channel 2(voltage from channel)
          '''
        
        s_state = self.app.sess_status
               
        if s_state.alive:
            try:
                self.app.data_lock.acquire()        # ACQUIRE DATA_LOCK <--
                skipSet = False
                # Test the latest value in the data array and update maximum if needed
                if (emglevel1 > s_state.peak[1]) or (emglevel2 > s_state.peak[1]):
                    if(emglevel1 > emglevel2):
                        new_max = emglevel1
                    else:
                        new_max = emglevel2
                        
                    # Checks to see if the incoming data is an abnormally high value    
                    if(new_max < self.app.cfg.MAX_Y_VALUE):
                        s_state.peak = (s_state.ch1_y_index, new_max)
                        self.redo_yauto = True
                    else:
                        skipSet = True
                        
                        
                # Check to see if new data screen is needed
                if(s_state.ch1_y_index > self.xaxis_max):
                    
                    self.renew_xaxis(emglevel1)
        
                if(not skipSet):
                    
                    jobs = self.app.job_queue
                    if self.start_time == 0:
                        self.start_time = time.time()
                    # Add the latest value to the data array and find current time
                    s_state.ch1_y_index = time.time() - self.start_time
                    self.ch1_y_data = np.append(self.ch1_y_data, emglevel1)
                    self.ch1_x_data = np.append(self.ch1_x_data, s_state.ch1_y_index)

    
                    # Add the latest value to the data array and find current time
                    s_state.ch2_y_index = time.time() - self.start_time
                    self.ch2_y_data = np.append(self.ch2_y_data, emglevel2)
                    self.ch2_x_data = np.append(self.ch2_x_data, s_state.ch2_y_index)

                    
                    # Add a draw() call to the job queue
                    jobs.put(self.draw)

            except:
                print "Exception in update_data()"
                print traceback.format_exc()
            finally:
                self.app.data_lock.release()        # RELEASE DATA_LOCK -->
            
    def draw_grid(self, boo):
        # Turn the plot grid on or off
        if boo:
            self.axes.grid(True, color='gray')
        else:
            self.axes.grid(False)
        

    def init_plot(self):
        """ Initial plot setup at the start of calibration and trials
        """
        self.reset_plot()
        
        alive = True       # begin adding jobs to the queue
        
        # Set the serial data acquisition thread to write values to arrays
        self.app.sess_status.alive = alive
                
    def set_fading(self):
        block = self.app.sess_status.block_index
        sch = self.app.cfg.fading_schedule[block]
        on = self.app.sess_status.plots_on
        off = self.app.sess_status.plots_off
        phase = self.app.sess_status.fading_phase
        if phase:
            if off < sch[1]:
                self.__set_fading(True)
                self.app.sess_status.plots_off += 1    # increment number of trials that have used fading (i.e. plot not visible)
            else:
                self.__set_fading(False)
                self.app.sess_status.plots_on += 1
                self.app.sess_status.plots_off = 0          # reset counter
                self.app.sess_status.fading_phase = False   # change fading phase
        else:
            if on < sch[0]:
                self.__set_fading(False)
                self.app.sess_status.plots_on += 1     # increment number of trials with plot visible
            else:
                self.__set_fading(True)
                self.app.sess_status.plots_off += 1
                self.app.sess_status.plots_on = 0           # reset counter
                self.app.sess_status.fading_phase = True    # change fading phase
        
    def __set_fading(self, boo):
        # Used to implement fading schedule. Should be called at the start of each plot
        # init_plot() > set_fading() > __set_fading()
        boo = not boo       # invert value to set plot visible is equal to inverse of fading on
        self.show_plot(boo)
        
    def show_plot(self, boo):
        # set plot visibility and update controls
        self.app.ctrl.chkbx_ctrl.cb_plot.SetValue(boo)
        self.emgplot.set_visible(boo)
        
    def reset_plot(self):
        # Clears the plot frame and resets data arrays
        self.app.sess_status.peak = (0, 0)
        self.app.sess_status.result = False
        self.getBackground = True
        
        # reset data arrays
        self.app.data_lock.acquire()      # ACQUIRE DATA_LOCK <--
        self.app.sess_status.y_index = Decimal(0)
        self.ch1_x_data = np.array([])
        self.ch1_y_data = np.array([])
        self.ch2_x_data = np.array([])
        self.ch2_y_data = np.array([])
        self.emgdata = []
        self.plot_state = []
        self.emgplot = self.axes.plot(self.ch1_x_data,
                                      self.ch1_y_data,
                                      'g',
                                      self.ch2_x_data,
                                      self.ch2_y_data,
                                      'b',
                                      animated=True)
        self.app.data_lock.release()      # RELEASE DATA_LOCK -->

        # reset the axes
        self.axes.set_xbound(lower=0, upper=self.app.cfg._x_axis_max)
        self.axes.set_ybound(lower=0, upper=self.app.cfg._y_axis_max)
        
        # Set grid lines
        self.draw_grid(self.plotctrl.cb_grid.GetValue())

        # Redraw the plot

        try:
            canvas = self.plotpanel.canvas
            self.background = canvas.copy_from_bbox(self.axes.bbox)
            
        except:
            print traceback.format_exc()
        self.Refresh()
        
        
    def draw(self):
        # Main draw function
        #    - Draws the changes in the model to the screen
        
        if self.app.sess_status == None:
            pass
        else:
            # Set y-axis limits
            # ONly if ymax button is pushed and an the limited need to be recalculated
            if (self.app.sess_status.ymax_auto and self.redo_yauto):
                if (self.app.sess_status.peak[1] > (self.app.cfg._y_axis_max * 0.8)) \
                    or (self.app.sess_status.peak[1] < (self.app.cfg._y_axis_max * 0.8)):
                    
                    self.app.cfg._y_axis_max = 1 + self.app.sess_status.peak[1] * (1 + (self.app.cfg.Y_MARGIN / 100.0))
                    self.axes.set_ybound(lower=0, upper=self.app.cfg._y_axis_max)
                    self.getBackground = True # Redo Visual Axis
                    self.redo_yauto = False
            
            try:        # main canvas update function call
                
                # Restore the background from the current plot values 
                if self.getBackground:
                    canvas = self.plotpanel.canvas
                    canvas.draw()
                    self.background = canvas.copy_from_bbox(self.axes.bbox)
                    self.getBackground = False
                    
                
                self.plotpanel.canvas.restore_region(self.background)
                # Set the plot data
                self.emgplot = self.axes.plot(self.ch1_x_data,
                                              self.ch1_y_data,
                                              'g',
                                              self.ch2_x_data,
                                              self.ch2_y_data,
                                              'b',
                                              animated=True)
                # Update the emg plot
                self.axes.draw_artist(self.emgplot[0])
                self.axes.draw_artist(self.emgplot[1])
                # Redraw axes rectangle
                self.plotpanel.canvas.blit(self.axes.bbox)
                
            except RuntimeError:
                print traceback.format_exc()
                print 'len x:', len(self.x_data), 'len y:', len(self.y_data)
                lenx = len(self.x_data)
                leny = len(self.y_data)
                diff = leny - lenx
                if diff > 0:        # if y_data has more items than x_data...
                    self.y_data = self.y_data[:-diff]   # remove items from y_data to even up the arrays
                elif diff < 0:      # if x_data has more items than y_data...
                    self.x_data = self.x_data[:-diff]   # remove items from x_data to even up the arrays
            except:
                print traceback.format_exc()
                    
    def target_visible(self, boo):
        self.target.set_visible(boo)
        
    def get_target(self):
        return self.target
    
    def get_axes(self):
        return self.axes
                    
        
# End of liveplot.py
